# -*- coding: utf-8 -*-
# 初始化加载数据库中图片向量
# @Time    : 2024/12/9 上午3:35
# @Author  : 赵鑫源
# @FileName: init_load_vectors.py
# @Software: PyCharm
import time

import numpy as np

# 全局变量
import global_vars
import my_logging
from mysql_helper import MySQLHelper
from tqdm import tqdm
from concurrent.futures import ThreadPoolExecutor
import threading

mysql_helper = MySQLHelper()  # 初始化
logging = my_logging.LogUtils().get_log()
global_lock = threading.Lock()  # 全局锁


def init_load_vectors(idx, row):
    """
        初始化加载数据库中的图片向量信息
        row : 数据库行数据
    """
    ulid = row['id']
    _vec = np.array([float(x) for x in row['imgVectors'].split(',')]).reshape(1, 512)
    # 使用锁保护共享资源
    with global_lock:
        global_vars.global_i2key[ulid] = row
        global_vars.global_index.add(_vec)
        global_vars.index_to_ulid[idx] = ulid  # 维护索引与 ULID 的映射


def get_img_info():
    """
        获取图片信息，并加载图片向量信息
    """
    rows = mysql_helper.execute_query(
        "select i.id,i.img_name,i.bucket_url,i.img_source_path,i.img_jpg_path,i.img_png_path,i.img_type,"
        "i.img_source_size,i.img_jpg_size,i.img_png_size,v.img_vectors from t_images_info i left join t_images_vectors v "
        "on i.id = v.img_id  where i.img_status = 'NORMAL' and v.img_vectors is not null ")
    fileSize = 0
    img_info_list = []

    # 使用线程池处理行数据
    with ThreadPoolExecutor(max_workers=10) as executor:
        # 处理行数据
        futures = [executor.submit(process_row, row) for row in rows]
        pbar = tqdm(total=len(rows), desc="处理图片数据", unit="条")
        for future in futures:
            result = future.result()
            if result:
                img_info_list.append(result)
                fileSize += 1
            pbar.update(1)
        pbar.close()

        # 加载图片向量
        futures = [executor.submit(init_load_vectors, idx, row) for idx, row in enumerate(img_info_list)]
        pbar = tqdm(total=len(img_info_list), desc="加载图片向量", unit="条")
        for future in futures:
            try:
                future.result()  # 等待任务完成
            except Exception as e:
                logging.error(f'加载图片向量失败: {e}')
            pbar.update(1)  # 更新进度条
        pbar.close()

    return fileSize





def process_row(row):
    """
        处理行数据，并返回图片信息
        row : 数据库行数据
    """
    return {
        'id': row['id'],
        'imgName': row['img_name'],
        'bucketUrl': row['bucket_url'],
        'imgSourcePath': row['img_source_path'],
        'imgJpgPath': row['img_jpg_path'],
        'imgPngPath': row['img_png_path'],
        'imgType': row['img_type'],
        'imgSourceSize': row['img_source_size'],
        'imgJpgSize': row['img_jpg_size'],
        'imgPngSize': row['img_png_size'],
        'imgVectors': row['img_vectors']
    }


class LoadImageVectors(object):
    def __init__(self):
        self.name = 'InitLoadImageVectors'
        count = get_img_info()
        logging.info(f'加载图片向量信息完成,共加载：{count}条')